import sys
import os
import json
import logging
import threading
from typing import Dict, Any, Optional
from contextlib import asynccontextmanager
from fastapi.middleware.cors import CORSMiddleware
# 第三方库
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain_core.runnables import RunnableConfig  # 新增：用于超时配置

# 解决导入路径问题：将项目根目录加入Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 本地模块导入
from model_chain.prompt.class_of_prompt.sqlTemplate import sqlPromptTemplate
from model_chain.prompt import prompt_config
from model_chain.prompt import creat_chat_chain as creat_chain
from model_chain.prompt import chat_history as chat_history
from model_chain.prompt.prompt_config import PromptConfig

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# 全局状态管理类 + 线程锁（解决并发安全问题）
class AppState:
    conversation_chain: Optional[Any] = None
    memory: Optional[Any] = None
    params: Optional[PromptConfig] = None
    memory_lock: threading.Lock = threading.Lock()  # 新增：记忆操作锁


app_state = AppState()


# 应用生命周期管理
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时初始化
    app_state.params = PromptConfig()
    # 创建默认对话链（使用sqlTemplate）
    init_prompt = sqlPromptTemplate().format(
        table_structures="",
        database_type="Mysql",
        output_example="",
        error_code=""
    )
    chain, memory = creat_chain.create_conversation_chain(init_prompt)
    app_state.conversation_chain = chain
    app_state.memory = memory
    logger.info("应用初始化完成")
    yield
    # 关闭时保存历史（加锁保护）
    with app_state.memory_lock:
        if app_state.memory:
            chat_history.save_conversation_history(app_state.memory)
    logger.info("应用关闭，历史对话已保存")


app = FastAPI(title="运维对话服务", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 开发环境允许所有来源（生产环境需指定具体IP/域名）
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法（GET/POST等）
    allow_headers=["*"],  # 允许所有请求头
)

# 请求/响应模型定义（保持不变）
class PromptClassRequest(BaseModel):
    class_of_prompt: int


class ConversationHistoryRequest(BaseModel):
    save_path: str


class ModelNameRequest(BaseModel):
    modelname: str


class FreeTemplateRequest(BaseModel):
    table_structures: str
    database_type: str = "Mysql"
    output_example: str = ""
    error_code: str = ""


class ApiResponse(BaseModel):
    response: int  # 0表示失败，1表示成功


def rebuild_chain(prompt: Optional[str] = None) -> bool:
    """
    重构对话链和记忆对象（增强版）
    确保在各种场景下都能生成有效链，返回True/False表示成功/失败
    """
    try:
        # 步骤1：参数校验与默认模板生成
        if not prompt:
            logger.info("未提供自定义模板，使用默认SQL模板")
            # 生成默认模板（与初始化逻辑一致，确保兼容性）
            try:
                prompt = sqlPromptTemplate().format(
                    table_structures="",
                    database_type="Mysql",
                    output_example="",
                    error_code=""
                )
                # 验证默认模板是否有效
                if not isinstance(prompt, str) or len(prompt.strip()) == 0:
                    raise ValueError("生成的默认模板为空或无效")
            except Exception as e:
                logger.error(f"默认模板生成失败: {str(e)}")
                # 使用极简 fallback 模板，确保不中断流程
                prompt = "你是SQL运维助手，请分析并解决问题。"

        # 步骤2：打印模板关键信息（脱敏，避免日志冗余）
        template_preview = prompt[:100] + "..." if len(prompt) > 100 else prompt
        # logger.info(f"使用模板重构对话链: {template_preview}")

        # 步骤3：加锁创建对话链（确保线程安全）
        with app_state.memory_lock:
            # 双重校验 create_conversation_chain 函数参数
            try:
                # 调用链创建函数，强制传入 prompt
                chain, memory = creat_chain.create_conversation_chain(free_prompt=prompt)
            except TypeError as e:
                # 处理参数名不匹配的兼容问题（如函数期望其他参数名）
                logger.warning(f"参数名不匹配，尝试兼容调用: {str(e)}")
                chain, memory = creat_chain.create_conversation_chain(prompt)  # 兼容旧参数名
            except Exception as e:
                logger.error(f"对话链创建失败: {str(e)}", exc_info=True)
                return False

            # 步骤4：验证链和记忆对象有效性
            if not chain or not memory:
                raise ValueError("create_conversation_chain 返回空的链或记忆对象")

            # 验证链是否可调用
            if not hasattr(chain, "invoke") and not hasattr(chain, "ainvoke"):
                raise TypeError("返回的链对象不支持 invoke/ainvoke 方法")

            # 步骤5：更新全局状态
            app_state.conversation_chain = chain
            app_state.memory = memory
        logger.info(prompt)
        logger.info("对话链重构成功")
        return True

    except Exception as e:
        # 捕获所有异常，确保返回False而非崩溃
        logger.error(f"对话链重构失败: {str(e)}", exc_info=True)
        # 失败时尝试恢复到上一个可用状态（如果有）
        if hasattr(app_state, "last_valid_chain") and hasattr(app_state, "last_valid_memory"):
            logger.info("尝试恢复到上一个可用状态")
            app_state.conversation_chain = app_state.last_valid_chain
            app_state.memory = app_state.last_valid_memory
        return False


# 以下接口逻辑不变，仅完善日志和异常处理
@app.post("/set-prompt-class", response_model=ApiResponse, status_code=200)
def set_prompt_class(request: PromptClassRequest):
    try:
        if not app_state.params:
            raise ValueError("配置未初始化")

        result = app_state.params.set_class_of_prompt(request.class_of_prompt)
        if not rebuild_chain():
            return {"response": 0}

        return {"response": 1 if result else 0}
    except Exception as e:
        logger.error(f"设置模板种类失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/set-conversation-history", response_model=ApiResponse, status_code=200)
def set_conversation_history(request: ConversationHistoryRequest):
    try:
        if not app_state.params:
            raise ValueError("配置未初始化")

        result = app_state.params.set_conversation_history(request.save_path)
        with app_state.memory_lock:  # 新增：加载历史时加锁
            if result and app_state.memory:
                chat_history.load_conversation_history(app_state.memory)

        if not rebuild_chain():
            return {"response": 0}

        return {"response": 1 if result else 0}
    except Exception as e:
        logger.error(f"设置对话历史失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/set-model", response_model=ApiResponse, status_code=200)
def set_model(request: ModelNameRequest):
    try:
        if not app_state.params:
            raise ValueError("配置未初始化")

        result = app_state.params.set_modelname(request.modelname)
        if not rebuild_chain():
            return {"response": 0}

        return {"response": 1 if result else 0}
    except Exception as e:
        logger.error(f"设置模型失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/set-free-template", response_model=ApiResponse, status_code=200)
def set_free_template(request: FreeTemplateRequest):
    try:
        free_prompt = sqlPromptTemplate().format(
            table_structures=request.table_structures,
            database_type=request.database_type,
            output_example=request.output_example,
            error_code=request.error_code
        )
        logger.info(free_prompt)
        if not rebuild_chain(free_prompt):
            logger.info("自定义模板设置失败")
            return {"response": 0}
        logger.info("自定义模板设置成功")
        return {"response": 1}
    except Exception as e:
        logger.error(f"设置自由模板失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(e))


# 重点优化：chat函数（异步化 + 超时控制 + 线程安全）
@app.post("/chat", response_model=Dict[str, Any])
async def chat(user_input: str):  # 改为异步函数
    """处理聊天请求（异步化改造，避免阻塞）"""
    if not app_state.conversation_chain or not app_state.memory:
        raise HTTPException(status_code=400, detail="对话链未初始化")

    try:
        # 处理特殊指令（加锁保存历史）
        if user_input.lower() == "退出":
            with app_state.memory_lock:
                chat_history.save_conversation_history(app_state.memory)
            return {
                "status": "exit",
                "message": "聊天结束，再见！",
                "response": 1
            }

        if user_input.lower() == "保存":
            with app_state.memory_lock:
                chat_history.save_conversation_history(app_state.memory)
            return {
                "status": "saved",
                "message": "对话已保存",
                "response": 1
            }

        # 核心优化1：异步调用模型（避免阻塞主线程）
        # 加载历史时加锁，防止并发读写冲突
        with app_state.memory_lock:
            chat_history_data = app_state.memory.load_memory_variables({})["chat_history"]

        # 核心优化2：设置超时（30秒），避免无限等待
        response = await app_state.conversation_chain.ainvoke(  # 异步调用ainvoke
            {"input": user_input, "chat_history": chat_history_data},
            config=RunnableConfig(timeout=300)  # 超时配置
        )

        # 核心优化3：加锁保存记忆，确保线程安全
        with app_state.memory_lock:
            response_str = json.dumps(response, ensure_ascii=False, indent=2)
            app_state.memory.save_context({"用户": user_input}, {"AI": response_str})
            history_count = len(app_state.memory.chat_memory.messages)

        return {
            "status": "success",
            "response": response,
            "history_count": history_count
        }
    except TimeoutError:
        logger.warning("模型调用超时")
        raise HTTPException(status_code=408, detail="模型响应超时，请稍后重试")
    except Exception as e:
        logger.error(f"对话处理失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/current-config", response_model=Dict[str, Any])
def get_current_config():
    try:
        if not app_state.params:
            raise ValueError("配置未初始化")

        with app_state.memory_lock:  # 加锁读取记忆
            history_count = len(app_state.memory.chat_memory.messages) if app_state.memory else 0

        return {
            "current_prompt_class": app_state.params.class_of_prompt,
            "current_model": app_state.params.modelname,
            "history_path": app_state.params.save_conversation_history,
            "max_history": app_state.params.max_history,
            "history_count": history_count
        }
    except Exception as e:
        logger.error(f"获取配置信息失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=str(e))


"""
启动命令（在项目根目录Chain_of_model下执行）:
uvicorn model_chain.run:app --host 0.0.0.0 --port 8000 --reload
"""